package com.weilus.filters;

import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.HttpMethod;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.PathContainer;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import java.util.*;

/**
 * spring:
 *   cloud:
 *     gateway:
 *       limiter:
 *       - id: flux-call
 *         limits:
 *         - GET /test/**=5/s           #5次每秒
 *         - /test/other=100/s          #100次每秒
 */
@Component
@Slf4j
public class TokenRateLimiterGatewayFilterFactory extends AbstractGatewayFilterFactory<TokenRateLimiterGatewayFilterFactory.Config> {

    private LimiterProperties properties;

    public TokenRateLimiterGatewayFilterFactory(LimiterProperties properties) {
        super(TokenRateLimiterGatewayFilterFactory.Config.class);
        this.properties = properties;
    }

    @Override
    public GatewayFilter apply(TokenRateLimiterGatewayFilterFactory.Config config) {
        return (exchange, chain) -> {
            PathContainer path = PathContainer.parsePath(exchange.getRequest().getURI().getRawPath());
            HttpMethod method = exchange.getRequest().getMethod();
            String token = Optional.ofNullable(exchange.getRequest().getHeaders().getFirst("token")).orElse(exchange.getRequest().getQueryParams().getFirst("token"));
            Route route =(Route) exchange.getAttributes().get(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
            if(!StringUtils.isEmpty(token)) {
                Set<RedisRateLimiter> matches = properties.getMatchers(route.getId(),method,path);
                if (!CollectionUtils.isEmpty(matches)) {
                    RedisRateLimiter limiter = matches.toArray(new RedisRateLimiter[]{})[0];
                    return limiter.isAllowed(null, token + path).flatMap((response) -> {
                        response.getHeaders().forEach((h, v) -> exchange.getResponse().getHeaders().add(h, v));
                        if (response.isAllowed()) {
                            return chain.filter(exchange);
                        } else {
                            log.debug("{} {} 限流 {}",method,path,limiter.getConfig());
                            ServerWebExchangeUtils.setResponseStatus(exchange, HttpStatus.TOO_MANY_REQUESTS);
                            return exchange.getResponse().setComplete();
                        }
                    });
                }
            }
            return chain.filter(exchange);
        };
    }

    public static class Config{

    }
}